#ifndef __UDP_PROXY_H__
#define __UDP_PROXY_H__
#include "stdafx.h"
#include <set>
#include <map>
#include <vector>
#include <mswsock.h>
#include <mstcpip.h>
#include "sync.h"
#include "socksdefs.h"
#include "iocp.h"
#include "threadpool.h"
#include "nfevents.h"
#include "proxydefs.h"

using namespace nfapi;

namespace UdpProxy
{
    class UDPProxyHandler
    {
    public:
        virtual void onUdpReceiveComplete(ENDPOINT_ID id, char *buf, int len, char *remoteAddress, int remoteAddressLen) = 0;
    };

    class UDPProxy : public IOCPHandler
    {
    public:
        UDPProxy();
        virtual ~UDPProxy();
        void *getExtensionFunction(SOCKET s, const GUID *which_fn);
        bool initExtensions();
        bool init(UDPProxyHandler *pProxyHandler,
                  char *proxyAddress, int proxyAddressLen,
                  const char *userName = NULL, const char *userPassword = NULL);
        void free();
        bool createProxyConnection(ENDPOINT_ID id);
        void deleteProxyConnection(ENDPOINT_ID id);
        bool udpSend(ENDPOINT_ID id, char *buf, int len, char *remoteAddress, int remoteAddressLen);
        OV_DATA *newOV_DATA();
        void deleteOV_DATA(OV_DATA *pov);
        bool startConnect(SOCKET socket, sockaddr *pAddr, int addrLen, ENDPOINT_ID id);
        bool startUdpReceive(SOCKET socket, ENDPOINT_ID id, OV_DATA *pov);
        bool startTcpReceive(SOCKET socket, ENDPOINT_ID id, OV_DATA *pov);
        bool startTcpSend(SOCKET socket, char *buf, int len, ENDPOINT_ID id);
        void onUdpSendComplete(SOCKET socket, DWORD dwTransferred, OV_DATA *pov, int error);
        void onUdpReceiveComplete(SOCKET socket, DWORD dwTransferred, OV_DATA *pov, int error);
        void setKeepAliveVals(SOCKET s);
        void onConnectComplete(SOCKET socket, DWORD dwTransferred, OV_DATA *pov, int error);
        void onTcpSendComplete(SOCKET socket, DWORD dwTransferred, OV_DATA *pov, int error);
        void onTcpReceiveComplete(SOCKET socket, DWORD dwTransferred, OV_DATA *pov, int error);
        virtual void onComplete(SOCKET socket, DWORD dwTransferred, OVERLAPPED *pOverlapped, int error);

    private:
        IOCPService m_service;
        UDPProxyHandler *m_pProxyHandler;

        tOvDataSet m_ovDataSet;
        tSocketMap m_socketMap;

        LPFN_CONNECTEX m_pConnectEx;

        char m_proxyAddress[NF_MAX_ADDRESS_LENGTH];
        int m_proxyAddressLen;

        std::string m_userName;
        std::string m_userPassword;

        AutoCriticalSection m_cs;
    };
}
struct UDP_CONTEXT
{
    UDP_CONTEXT(PNF_UDP_OPTIONS options)
    {
        if (options)
        {
            m_options = (PNF_UDP_OPTIONS) new char[sizeof(NF_UDP_OPTIONS) + options->optionsLength];
            memcpy(m_options, options, sizeof(NF_UDP_OPTIONS) + options->optionsLength - 1);
        }
        else
        {
            m_options = NULL;
        }
    }

    ~UDP_CONTEXT()
    {
        if (m_options)
            delete[] m_options;
    }

    PNF_UDP_OPTIONS m_options;
};
#endif // !__UDP_PROXY_H__